
package com.intel.hibench.flinkbench.datasource;

import com.intel.hibench.flinkbench.util.KeyedTupleSchema;
import com.intel.hibench.flinkbench.util.FlinkBenchConfig;

import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;

import java.io.IOException;
import java.util.Properties;

public abstract class StreamBase {

  private SourceFunction<Tuple2<String, String>> dataStream;

  public SourceFunction<Tuple2<String, String>> getDataStream() {
    return this.dataStream;
  }

  public void createDataStream(FlinkBenchConfig config) throws Exception {

    Properties properties = new Properties();
    properties.setProperty("zookeeper.connect", config.zkHost);
    properties.setProperty("group.id", config.consumerGroup);
    properties.setProperty("bootstrap.servers", config.brokerList);
    properties.setProperty("auto.offset.reset", config.offsetReset);

    this.dataStream = new FlinkKafkaConsumer<Tuple2<String, String>>(
            config.topic,
            new TupleDeserializationSchema(),
            properties);
  }

  public void processStream(FlinkBenchConfig config) throws Exception {
  }

  private static class TupleDeserializationSchema implements DeserializationSchema<Tuple2<String, String>> {

    @Override
    public Tuple2<String, String> deserialize(byte[] message) throws IOException {
      // 假设消息格式为 "key,value"
      String messageStr = new String(message);
      String[] parts = messageStr.split(",");
      if (parts.length != 2) {
        throw new IOException("Message format is incorrect: " + messageStr);
      }
      return new Tuple2<String, String>(parts[0], parts[1]);
    }

    @Override
    public boolean isEndOfStream(Tuple2<String, String> nextElement) {
      return false;
    }

    @Override
    public TypeInformation<Tuple2<String, String>> getProducedType() {
      return TypeInformation.of(new TypeHint<Tuple2<String, String>>() {});
    }
  }
}